# 加载 li.txt
from langchain.document_loaders import TextLoader
loader = TextLoader("li.txt")
documents = loader.load()


# 文档分割
from langchain.text_splitter import CharacterTextSplitter
# 创建拆分器
text_splitter = CharacterTextSplitter(chunk_size=128, chunk_overlap=0)
# 拆分文档
documents = text_splitter.split_documents(documents)

# 接下来对分割后的数据进行embedding，并写入数据库。这里选用
# m3e-base作为embedding模型，向量数据库选用Chroma

from langchain.embeddings import OllamaEmbeddings
from langchain.vectorstores import Chroma

# 使用 Ollama 的 nomic-embed-text 模型
embedding = OllamaEmbeddings(
    model="nomic-embed-text",
    base_url="http://localhost:11434"
)

# load data to Chroma db
db = Chroma.from_documents(documents, embedding)
# similarity search
db.similarity_search("藜一般在几月播种？")

from langchain.llms import Ollama
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

# 初始化 Ollama LLM
llm = Ollama(
    model="llama3.3:70b",
    base_url="http://localhost:11434"
)

# 创建自定义提示模板
prompt_template = """基于以下已知信息，简洁和专业地回答用户的问题。
如果无法从中得到答案，请说 "根据已知信息无法回答该问题"。

已知信息：
{context}

问题：{question}
回答："""

PROMPT = PromptTemplate(
    template=prompt_template,
    input_variables=["context", "question"]
)

# 创建检索问答链
retriever = db.as_retriever(search_kwargs={"k": 3})
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    return_source_documents=True,
    chain_type_kwargs={"prompt": PROMPT}
)

# 进行问答
response = qa_chain({"query": "藜怎么防治虫害？"})
print(response['result'])

